import pandas as pd
import numpy as np
import jieba
import re
from transformers import BertTokenizer, BertModel
import torch
import lightgbm as lgb
import json
import joblib
from sklearn.metrics import accuracy_score
import math
from  sklearn.model_selection import GridSearchCV
from sklearn.model_selection import train_test_split

# Check if a GPU is available
if torch.cuda.is_available():
    # If a GPU is available, move the input tensors and the model to the GPU
    device = torch.device("cuda")
else:
    device = torch.device("cpu")

# Load the BERT model and tokenizer
model = BertModel.from_pretrained('bert-base-chinese').to(device)
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')

# Load the input data
df4=pd.read_csv('./predata_q.csv',encoding='utf-8')
q = df4["q"].values[0:1500]
content = df4["content"].values[0:1500]

q1=list(df4['qidx'].values)[0:1500]
q2=list(df4['didx'].apply(str).values)[0:1500]

pred_all=[]
for i in range(len(q1)):
    pred_all.append([q1[i],q2[i]])
  
# Encode the input data using BERT and concatenate the resulting vectors
train = []
for i in range(len(q1)):
    query_tokens = tokenizer.encode(q[i],max_length=512,truncation=True,return_tensors='pt')
    query_vec = model(query_tokens.to(device))[0][:,0,:].detach().cpu().numpy()
    candidate_tokens = tokenizer.encode(content[i],max_length=512,truncation=True,return_tensors='pt')
    candidate_vec = model(candidate_tokens.to(device))[0][:,0,:].detach().cpu().numpy()
    vec = np.concatenate((query_vec, candidate_vec), axis=0).reshape(-1)
    train.append(vec)
train = np.array(train)

# Load the labels and train the LightGBM classifier
label = df4["label"].values[0:1500]

X_train, X_test, Y_train, Y_test = train_test_split(train,label,random_state=12,stratify=label,test_size=0.3,shuffle=True)


# gbm=lgb.LGBMClassifier(num_leaves=7,
#                        reg_alpha=1,
#                        reg_lambda=0.1,
#                        max_depth=3,
#                        objective='lambdarank',
#                        learning_rate=0.1,
#                        min_child_samples=20,
#                        n_estimators=5000,
#                        subsample=0.8,
#                        colsample_bytree=0.8,
#                        random_state=9487)

#step1 调整max_depth 和 num_leaves
parameters = {
    'max_depth': [2,4,6],
    'num_leaves': [10,20,30],
}

gbm=lgb.LGBMClassifier(num_leaves=10,
                       reg_alpha=1,
                       reg_lambda=0.1,
                       max_depth=2,
                       objective='lambdarank',
                       learning_rate=0.1,
                       num_iterations = 200,
                       min_child_samples=20,
                       n_estimators=5000,
                       subsample=0.8,
                       colsample_bytree=0.8,
                       random_state=9487)

gbm.fit(train,label)


joblib.dump(gbm, './loan_model3.pkl')

pre=gbm.predict(train)


for i in range(len(pred_all)):
    pred_all[i].append(pre[i])

pred_all=sorted(pred_all,key=lambda x:x[2],reverse=True)

dic={}
for x,y,z in pred_all:
    if str(x) not in dic:
        dic[str(x)]=[]
    dic[str(x)].append(y)

def ndcg(ranks, K):
    dcg_value = 0.
    idcg_value = 0.

    sranks = sorted(ranks, reverse=True)

    for i in range(0,K):
        logi = math.log(i+2,2)
        dcg_value += ranks[i] / logi
        idcg_value += sranks[i] / logi
    if idcg_value == 0.0:
        idcg_value += 0.00000001
    return dcg_value/idcg_value

def cal_ndcg(all_preds, all_labels):
    ndcgs = []
    for qidx, pred_ids in all_preds.items():
        did2rel = all_labels[qidx]
        ranks = [did2rel[idx] if idx in did2rel else 0 for idx in pred_ids]
        ndcgs.append(ndcg(ranks, 30))
        print(f'********** qidx: {qidx} **********')
        print(f'top30 pred_ids: {pred_ids}')
        print(f'ranks: {ranks}')
    print(ndcgs)
    return sum(ndcgs) / len(ndcgs)




json.dump(dic, open('./prediction3.json', "w", encoding="utf8"), indent=2,ensure_ascii=False)
# Evaluate the model
accuracy = accuracy_score(label,pre)
print("Accuracy:", accuracy)

all_label={"259": {"11729": 1, "29495": 0, "35043": 3, "21982": 0, "40253": 0, "16260": 2, "39097": 1, "40181": 2, "15356": 0, "28966": 0, "2204": 0, "9216": 3, "22447": 3, "19078": 2, "11290": 3, "15019": 0, "30042": 2, "41744": 2, "3775": 3, "9316": 2, "26173": 0, "10475": 0, "11423": 0, "37816": 0, "21306": 3, "14402": 0, "721": 0, "41600": 2, "42756": 1, "34160": 2}, "221": {"41364": 2, "9238": 3, "678": 2, "7719": 3, "33425": 3, "37093": 2, "16132": 2, "2151": 3, "2768": 2, "13896": 3, "11164": 3, "11527": 3, "43700": 3, "26009": 2, "20531": 3, "23178": 3, "5956": 2, "39221": 2, "35829": 2, "28250": 3, "12432": 3, "7194": 2, "17650": 3, "15308": 2, "705": 3, "13179": 3, "11871": 2, "23638": 3, "28399": 3, "42863": 3}, "2132": {"3940": 1, "7601": 3, "21384": 3, "40445": 3, "24527": 1, "13198": 3, "16385": 2, "14770": 1, "32543": 2, "11128": 2, "28377": 1, "39960": 1, "25911": 1, "37697": 2, "36141": 3, "19028": 2, "35124": 1, "24360": 2, "18469": 2, "9752": 2, "34387": 1, "32596": 3, "32311": 2, "4263": 2, "26388": 3, "28808": 2, "3779": 1, "4111": 1, "36183": 2, "34444": 1}, "2143": {"14314": 3, "12420": 2, "929": 1, "21257": 3, "24048": 2, "11079": 3, "28066": 3, "26125": 2, "31717": 2, "7592": 2, "17620": 2, "6683": 2, "32293": 0, "37077": 2, "4421": 1, "10349": 2, "376": 0, "35137": 0, "11642": 2, "39827": 3, "41416": 0, "42234": 2, "43587": 0, "40157": 2, "7891": 1, "22777": 1, "42644": 0, "42645": 0, "30930": 1, "31465": 1}, "1972": {"13549": 0, "35634": 1, "14825": 3, "18143": 2, "16083": 2, "426": 2, "9720": 2, "17725": 2, "2597": 2, "6683": 2, "38405": 2, "7030": 2, "38406": 2, "9752": 2, "6155": 2, "27790": 2, "43345": 3, "31873": 2, "32882": 2, "28661": 2, "4046": 2, "25799": 2, "22604": 3, "34725": 3, "7498": 2, "23980": 3, "19327": 3, "37697": 2, "41707": 2, "11056": 2}, "1978": {"18291": 1, "20533": 1, "29730": 1, "31085": 2, "11382": 2, "25367": 1, "5481": 0, "20632": 1, "33545": 1, "14648": 3, "33342": 1, "33117": 2, "18599": 1, "10402": 3, "32142": 0, "1424": 1, "21982": 2, "26493": 1, "35043": 1, "16310": 1, "25040": 2, "17948": 1, "8412": 1, "12510": 2, "21530": 2, "14825": 1, "5230": 1, "21431": 2, "5310": 2, "26423": 3}, "2361": {"1671": 2, "22403": 2, "41696": 3, "41695": 3, "38119": 2, "22477": 2, "30904": 2, "37643": 1, "25628": 2, "19049": 0, "30583": 3, "24817": 2, "19592": 2, "22459": 2, "34486": 2, "17848": 1, "10475": 2, "28504": 2, "31358": 2, "21705": 2, "4068": 2, "11239": 2, "43756": 3, "27534": 2, "40908": 2, "504": 2, "42379": 2, "18422": 1, "26627": 2, "10014": 2}, "2373": {"30548": 2, "35919": 2, "10845": 2, "19828": 3, "37736": 3, "37737": 3, "32351": 3, "4935": 1, "27406": 2, "38429": 2, "28504": 2, "24208": 3, "25938": 3, "23849": 3, "37995": 2, "22713": 3, "22186": 2, "7760": 2, "40344": 1, "36844": 2, "449": 3, "2818": 3, "9153": 1, "42102": 3, "15414": 2, "17061": 1, "28862": 3, "25561": 2, "18287": 2, "10731": 3}, "2331": {"43328": 2, "30491": 2, "32859": 3, "43330": 2, "14169": 3, "23275": 2, "16478": 2, "30723": 2, "11977": 3, "5618": 3, "42249": 2, "43365": 3, "16611": 3, "30369": 2, "16547": 2, "19301": 2, "39326": 1, "14520": 2, "15552": 1, "43767": 3, "9096": 2, "11693": 3, "15451": 3, "35460": 2, "28504": 2, "18287": 3, "1985": 2, "32442": 2, "39460": 2, "2935": 2}, "3228": {"18406": 3, "15836": 3, "22657": 3, "37269": 3, "26016": 3, "32484": 3, "25240": 3, "30689": 3, "120": 2, "482": 3, "18066": 3, "21678": 1, "27563": 3, "14800": 2, "42997": 3, "1070": 3, "24111": 3, "38283": 3, "3631": 3, "291": 3, "29329": 1, "15054": 3, "32643": 2, "13049": 3, "34018": 3, "8964": 3, "35447": 2, "5336": 3, "13888": 3, "428": 3}, "3746": {"10660": 3, "16090": 2, "30278": 3, "25589": 2, "34121": 3, "4760": 2, "29849": 2, "3511": 2, "18342": 3, "19237": 2, "40590": 2, "38951": 2, "36766": 2, "22959": 3, "10346": 2, "25368": 2, "12543": 2, "6917": 2, "43639": 2, "36076": 1, "23876": 3, "27500": 2, "38269": 3, "38099": 3, "41670": 2, "19782": 2, "38100": 3, "288": 2, "1707": 3, "526": 2}, "3765": {"11013": 3, "19252": 2, "20066": 2, "29449": 2, "217": 2, "30733": 3, "5042": 3, "36806": 3, "24669": 3, "32011": 3, "3150": 3, "23393": 2, "32027": 2, "16218": 2, "20593": 2, "19533": 2, "25794": 2, "31050": 2, "521": 2, "9568": 3, "16148": 2, "16312": 2, "11924": 2, "19111": 2, "2590": 3, "40507": 2, "37005": 2, "22997": 2, "25368": 2, "5997": 2}, "3342": {"26458": 3, "2065": 3, "23089": 3, "2183": 3, "23533": 3, "2339": 3, "23001": 3, "3472": 3, "16554": 3, "12546": 3, "29812": 3, "26912": 3, "26987": 3, "10633": 3, "202": 3, "35316": 3, "37369": 3, "22464": 3, "33579": 3, "22748": 3, "31404": 3, "19341": 3, "26692": 3, "31263": 3, "7805": 3, "19295": 3, "2587": 3, "12253": 3, "33983": 3, "5740": 3}, "1405": {"32697": 3, "41447": 3, "9189": 2, "25460": 2, "25817": 2, "36054": 2, "36083": 2, "22818": 3, "8587": 3, "43382": 2, "27689": 2, "15412": 2, "31714": 2, "974": 3, "10115": 3, "40819": 2, "10920": 2, "14341": 2, "23054": 3, "16411": 2, "35328": 2, "36119": 2, "1969": 2, "17959": 2, "13559": 2, "36309": 2, "28261": 3, "28691": 2, "7121": 2, "40491": 2}, "1430": {"6304": 2, "34101": 2, "2596": 2, "1054": 2, "5577": 2, "32910": 1, "33250": 1, "31960": 3, "36700": 1, "7964": 2, "25210": 1, "4401": 2, "21090": 1, "34968": 2, "43519": 1, "23298": 2, "2517": 2, "29182": 2, "2731": 2, "16075": 1, "31347": 1, "30820": 2, "17031": 1, "37656": 2, "38370": 2, "16867": 1, "24399": 1, "2987": 2, "13662": 2, "17105": 2}, "1325": {"8201": 2, "25005": 2, "25366": 1, "10154": 2, "8969": 3, "42142": 2, "21958": 0, "36503": 0, "16383": 1, "36212": 0, "30513": 1, "18815": 1, "7861": 1, "4119": 0, "7419": 0, "6552": 1, "25244": 1, "10068": 1, "43006": 1, "16230": 1, "608": 2, "34642": 1, "29623": 1, "29476": 0, "12373": 0, "872": 1, "23030": 1, "34302": 1, "28776": 0, "12977": 2}, "1355": {"36109": 1, "37160": 1, "28436": 0, "5894": 0, "4875": 0, "16713": 0, "18333": 0, "13051": 0, "35862": 1, "3881": 0, "42967": 0, "14367": 1, "1120": 0, "22522": 0, "5055": 0, "13524": 1, "32888": 0, "41481": 0, "4649": 1, "41480": 0, "25453": 1, "32096": 0, "38347": 1, "6687": 1, "39259": 1, "7142": 1, "41482": 1, "15605": 0, "5987": 0, "10158": 2, "50004": 3, "50005": 3}, "4738": {"33955": 0, "13295": 2, "7453": 2, "40643": 2, "31071": 2, "30032": 2, "8650": 0, "23980": 2, "25493": 2, "33652": 2, "23954": 2, "21167": 2, "40342": 2, "11163": 0, "34292": 2, "574": 2, "8582": 2, "36611": 2, "11042": 1, "8143": 1, "30692": 2, "31090": 2, "32097": 2, "3190": 2, "39028": 1, "2797": 2, "33630": 2, "26254": 2, "40343": 2, "23164": 2, "50006": 3, "50007": 3, "50008": 3}, "4794": {"17367": 3, "42599": 3, "11007": 2, "42600": 3, "14825": 3, "24603": 2, "6155": 3, "20673": 2, "1804": 3, "7670": 3, "38906": 3, "6683": 2, "38297": 3, "30176": 3, "38296": 2, "1634": 3, "23980": 3, "16573": 3, "41980": 3, "29334": 3, "39263": 3, "34662": 2, "39262": 3, "10811": 3, "38905": 3, "34536": 2, "34725": 3, "19359": 3, "3398": 3, "40036": 2}, "4829": {"30009": 2, "25242": 2, "40795": 3, "1634": 2, "30032": 3, "25987": 2, "7724": 3, "41337": 2, "9189": 1, "31883": 2, "11232": 1, "35003": 2, "28405": 2, "23998": 2, "34163": 2, "21670": 3, "34292": 2, "38296": 1, "34761": 3, "38297": 2, "31643": 2, "1837": 2, "34222": 2, "9339": 2, "29327": 3, "4046": 1, "13921": 2, "27523": 3, "32237": 2, "42610": 1}}
vv=cal_ndcg(dic,all_label)
print(vv)

#2023.1.6模型调参
gsearch = GridSearchCV(gbm, param_grid=parameters, scoring='average_precision', cv=3)
gsearch.fit(X_train, Y_train)
print('参数的最佳取值:{0}'.format(gsearch.best_params_))
print('最佳模型得分:{0}'.format(gsearch.best_score_))
print(gsearch.cv_results_['mean_test_score'])
print(gsearch.cv_results_['params'])